import requests
import json
from datetime import datetime, timedelta, timezone
from astropy.coordinates import SkyCoord, AltAz, get_body, EarthLocation, get_moon
from astropy.time import Time
import numpy as np
import os
from astropy import units as u
from scipy.optimize import minimize
from parseGWACData import getGWACSkyListByDate
from parseSVOMData import getProcessedSvomWorkplan, compute_overlap_xinglong, compute_svom_overlay_local_zenith
from dataDraw import draw_overlap_ratios, plot_overlap_xinglong
import tools
import random
import skyFieldPlot as skyPlot
import sys
import shutil

# 观测站地点（中国科学院国家天文台兴隆观测站）
# （经度 117.575°，纬度 40.393°，海拔 960m）
obs_location = EarthLocation.from_geodetic(117.575, 40.393, 960)  # 经纬度和海拔（米）


def parse_gwac_data(gwac_data):
    gwac_observations = []
    for entry in gwac_data:
        time = datetime.fromisoformat(entry["mintime"].replace("Z", "")) + timedelta(hours=8)
        sky_ids = entry["sky_ids_json"]
        telescope_data = {}
        
        field_ids = [] #过滤重复天区
        for sky_id in sky_ids:
            telescope_id, field_id = sky_id.split("-")
            if field_id not in field_ids:
                field_ids.append(field_id)
            else:
                continue
            
            ra = float(field_id[:3]) + float(field_id[3]) / 10
            dec = float(field_id[5:7]) + float(field_id[7]) / 10
            if field_id[4] == "1":
                dec = -dec
            
            telescope_data[telescope_id] = SkyCoord(ra, dec, unit="deg")
        
        workingMountNum = len(sky_ids)
        validSkyNum = len(field_ids)
        gwac_observations.append({
            'time': time,
            'workingMountNum': workingMountNum,
            'validSkyNum': validSkyNum,
            'telescope_data': telescope_data
        })
    return gwac_observations

def parse_svom_data(svom_data):
    observations = []
    for obs in svom_data:
        category = obs['category']
        start_time = obs['start']
        end_time = obs['end']
        center = SkyCoord(obs["ra"], obs["dec"], unit="deg")
        observations.append([start_time, end_time, center, category])
    return observations

'''
使用线性比例估算重叠面积
'''
def compute_overlap2(gwac_obs, svom_obs):
    
    gwac_field_radius = 9.3  #单位度， GWAC 视场半径
    svom_field_radius = 90  # 单位度， SVOM 观测区域半径
    gwac_field_area = 18.6 * 18.6  # 每个GWAC视场的面积
    max_distance = gwac_field_radius + svom_field_radius
    
    all_overlap_areas = []

    for tgObs in gwac_obs:
        time = tgObs['time']
        telescopes = tgObs['telescope_data']
        workingMountNum = tgObs['workingMountNum']
        validSkyNum = tgObs['validSkyNum']
        
        if time.hour < 9 or time.hour >= 23:
            continue
        
        total_overlap_area = 0
        overlap_areas = []
        
        tsvomObs = []
        for _, gwac_field in telescopes.items():
            tmps = []
            if len(tsvomObs)==0:
                for tobs in svom_obs:
                    start, end, center, category = tobs
                    if start <= time <= end:
                        tsvomObs = tobs
                        break
            
            if len(tsvomObs)==4:
                start, end, center, category = tsvomObs
                separation = gwac_field.separation(center).degree
                if separation < max_distance: 
                    if separation < svom_field_radius-gwac_field_radius: # GWAC 视场完全在 SVOM 内
                        overlap_fraction = 1
                    else:
                        overlap_fraction = (max_distance - separation.item()) / 18.6  # 计算重叠比例
                    total_overlap_area += overlap_fraction
                    overlap_areas.append(overlap_fraction)
        
        if len(tsvomObs)==4:
            svom_obs_type = tsvomObs[3]
        else:
            svom_obs_type = ''
        all_overlap_areas.append({
            'time': time.isoformat(),
            'total_overlap_area_valid_sky': total_overlap_area / validSkyNum,
            'total_overlap_area_10mount_sky': total_overlap_area / 10,
            'overlap_areas': overlap_areas,
            'workingMountNum': workingMountNum,
            'validSkyNum': validSkyNum,
            'svom_obs_type': svom_obs_type
        })
    
    return all_overlap_areas

'''
使用球面公式计算重叠面积
'''
def compute_overlap3(gwac_obs, svom_obs):
    gwac_field_radius = 9.3 * u.deg  # GWAC 视场半径
    svom_field_radius = 90 * u.deg   # SVOM 视场半径
    max_distance = gwac_field_radius + svom_field_radius

    # 计算 GWAC 视场面积
    gwac_field_area = np.pi * (gwac_field_radius ** 2)

    all_overlap_areas = []

    for tgObs in gwac_obs:
        time = tgObs['time']
        telescopes = tgObs['telescope_data']
        workingMountNum = tgObs['workingMountNum']
        validSkyNum = tgObs['validSkyNum']
        
        # 过滤时间范围
        if time.hour < 9 or time.hour >= 23:
            continue
        
        # 查找当前时间的 SVOM 观测
        tsvomObs = None
        for tobs in svom_obs:
            start, end, center, category = tobs
            if start <= time <= end:
                tsvomObs = tobs
                break
        
        if tsvomObs is None:
            continue
        
        start, end, svom_center, svom_category = tsvomObs
        
        localZenithOverlap30 = compute_svom_overlay_local_zenith(svom_center, time, obs_location)
        
        total_overlap_area = 0 * u.deg**2  # 初始化 total_overlap_area 并赋予单位
        overlap_areas = []
        
        # 计算每个 GWAC 望远镜与 SVOM 的重叠
        for _, gwac_field in telescopes.items():
            separation = gwac_field.separation(svom_center)
            if separation > max_distance:
                continue  # 无重叠
            
            if separation < svom_field_radius - gwac_field_radius:
                # GWAC 视场完全在 SVOM 内
                overlap_area = np.pi * (gwac_field_radius ** 2)
            else:
                # 计算重叠面积
                d = separation.to(u.deg).value
                r1 = gwac_field_radius.to(u.deg).value
                r2 = svom_field_radius.to(u.deg).value
                overlap_area = (
                    r1**2 * np.arccos((d**2 + r1**2 - r2**2) / (2 * d * r1)) +
                    r2**2 * np.arccos((d**2 + r2**2 - r1**2) / (2 * d * r2)) -
                    0.5 * np.sqrt((-d + r1 + r2) * (d + r1 - r2) * (d - r1 + r2) * (d + r1 + r2)))
                overlap_area *= (np.pi / 180) ** 2  # 转换为立体弧度
                overlap_area *= u.deg**2  # 添加单位
            
            total_overlap_area += overlap_area
            overlap_areas.append(overlap_area)
        
        # 所有转台重叠区域面积之和（相对于GWAC视场）
        overlap_ratio = (total_overlap_area / gwac_field_area).value  # 去掉单位
        
        # 计算相对于 validSkyNum 和 10 的重叠比例
        overlap_ratio_valid_sky = overlap_ratio / validSkyNum #有效的天区数作为分母
        overlap_ratio_10mount_sky = overlap_ratio / 10 #所有转台数作为分母
        
        all_overlap_areas.append({
            'time': time.isoformat(),
            'overlap_ratio': overlap_ratio,  # 所有转台重叠区域面积之和 相对于 GWAC 视场的重叠比例
            'total_overlap_area_valid_sky': overlap_ratio_valid_sky, 
            'total_overlap_area_10mount_sky': overlap_ratio_10mount_sky, #
            # 'overlap_areas': [area.value for area in overlap_areas],  # 去掉单位
            'workingMountNum': workingMountNum,
            'validSkyNum': validSkyNum,
            'svom_obs_type': svom_category,
            'localZenithOverlap30': localZenithOverlap30
        })
    
    return all_overlap_areas

    
def compute_alt_az(time, coord):
    astropy_time = Time(time)
    altaz = coord.transform_to(AltAz(obstime=astropy_time, location=obs_location))
    return altaz.alt.degree, altaz.az.degree

def compute_telescope_angles(gwac_obs):
    telescope_angles = {}
    
    # for time, telescopes in gwac_obs:
    for tgObs in gwac_obs:
        time = tgObs['time']
        telescopes = tgObs['telescope_data']
        if time.hour < 9 or time.hour >= 23:
            continue
        
        for telescope_id, coord in telescopes.items():
            alt, az = compute_alt_az(time, coord)
            if telescope_id not in telescope_angles:
                telescope_angles[telescope_id] = []
            telescope_angles[telescope_id].append((time.isoformat(), alt, az))
    
    return telescope_angles

def compute_moon_altitudes(gwac_obs):
    moon_altitudes = []
    
    # for time, _ in gwac_obs:
    for tgObs in gwac_obs:
        time = tgObs['time']
        # telescopes = tgObs['telescope_data']
        if time.hour < 9 or time.hour >= 23:
            continue
        
        # moon_coord = get_moon(Time(time), location=obs_location)
        moon_coord = get_body("moon", Time(time), location=obs_location)
        alt, _ = compute_alt_az(time, moon_coord)
        moon_altitudes.append((time.isoformat(), alt))
    
    return moon_altitudes

def json_serial(obj):
    if isinstance(obj, datetime):
        return obj.isoformat()
    
def prepareGWACSVOMData(dateStr=None):
    
    if dateStr is not None:
        utcTimeStr = dateStr
    else:
        utcTimeStr = tools.getCurDateStr()
    print(f"start processing {utcTimeStr}")

    # utcTimeStr = "2025-02-27"
    saveDir = "public/svomSummary/data/"+utcTimeStr
    if not os.path.exists(saveDir):
        os.mkdir(saveDir)
    
    if dateStr is not None:
        gwac_data = getGWACSkyListByDate(saveDir, dateStr)
    else:
        gwac_data = getGWACSkyListByDate(saveDir)
        
    if len(gwac_data)>30: #当天观测时间超过30分钟，gwac_data中每条记录1分钟
        
        oneDaySummaryFile = saveDir + "/oneDaySummary.json"
        if os.path.exists(oneDaySummaryFile):
            print(f"{utcTimeStr} has been processed before, skip it.")
            return
        
        svom_data = getProcessedSvomWorkplan(saveDir, utcTimeStr)
        
        gwac_obs = parse_gwac_data(gwac_data)
        svom_obs = parse_svom_data(svom_data)
        
        #compute_overlap2 和compute_overlap3 的结果基本一致，但是compute_overlap3 的结果更准确
        # overlap_ratios = compute_overlap2(gwac_obs, svom_obs)
        overlap_ratios = compute_overlap3(gwac_obs, svom_obs)
        # telescope_angles = compute_telescope_angles(gwac_obs)
        # moon_altitudes = compute_moon_altitudes(gwac_obs)
        telescope_angles = []
        moon_altitudes = []
        
        gwacSummary = {
            "overlap_ratios": overlap_ratios,
            "telescope_angles": telescope_angles,
            "moon_altitudes": moon_altitudes,
            "svom_workplans": svom_data
        }
        
        rstSavePath = saveDir + "/observation_results.json"
        with open(rstSavePath, "w") as f:
            # json.dump(gwacSummary, f, indent=4)
            json.dump(gwacSummary, f, default=json_serial)

        saveImagePath = f"{saveDir}/overlap_ratios.png"
        # 绘制四幅图：GWAC的10个转台综合视场在SVOM的覆盖比例，地平30度以上天区在SVOM的覆盖比例，GWAC有效转台数目，SVOM workplan
        draw_overlap_ratios(rstSavePath, saveImagePath)
        
        saveImagePaths = skyPlot.batch_plot_sky_coverage(gwac_obs, svom_obs, overlap_ratios, saveDir)
        
        oneDaySummary = {
            "gwac_data_file": "gwac_data.json",
            "svom_data_file": "svom_data.json",
            "overlap_ratios_file": "observation_results.json",
            "overlap_ratios_image": "overlap_ratios.png",
            "sky_coverage_images": saveImagePaths
        }
        
        with open(oneDaySummaryFile, "w") as f:
            json.dump(oneDaySummary, f, indent=4)
            
    else:
        print(f"obsDay {utcTimeStr} have observation less than 30 minutes, remove data dir {saveDir}")
        if os.path.exists(saveDir):
            # 删除目录及其所有内容
            shutil.rmtree(saveDir)
    
def test2():
    
    # utcTimeStr = tools.getCurDateStr()
    utcTimeStr = "2025-02-26"
    saveDir = "public/svomSummary/data/"+utcTimeStr
    if not os.path.exists(saveDir):
        os.mkdir(saveDir)
    svom_data = getProcessedSvomWorkplan(saveDir)
    svom_obs = parse_svom_data(svom_data)
    overlap_areas = compute_overlap_xinglong(svom_obs)
    plot_overlap_xinglong(overlap_areas)
    
if __name__ == "__main__":
    
    if len(sys.argv) == 2:
        #"2025-02-27"
        dateStr = sys.argv[1]
        prepareGWACSVOMData(dateStr)
    else:
        prepareGWACSVOMData()
    
